"""
数据整合模块 - 负责将CSV数据转换为评估工具所需的JSON格式，并调用评估流程
"""
import csv
import json
import os
import subprocess
from typing import List, Dict, Any, Optional, Tuple
import pandas as pd

def convert_csv_to_json(evaluation_csv: str, 
                       reranking_csv: str, 
                       output_dir: str,
                       question_id_column: str = "问题"):
    """
    将CSV数据转换为评估工具所需的JSON格式
    
    参数:
        evaluation_csv: 评估结果CSV文件路径
        reranking_csv: 重排结果CSV文件路径
        output_dir: 输出目录
        question_id_column: 问题列名（用于分组）
    
    返回:
        生成的JSON文件路径的字典
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 读取CSV文件
    eval_df = pd.read_csv(evaluation_csv)
    rerank_df = pd.read_csv(reranking_csv)
    
    # 1. 准备评估数据JSON (包含问题、知识片段和黄金标准得分)
    eval_data = {}
    
    # 按问题ID分组
    grouped = eval_df.groupby(question_id_column)
    for question_id, group in grouped:
        passages = group["知识块"].tolist()
        gold_scores = group["评分"].tolist()
        
        # 获取问题文本 (假设问题列存在)
        question_text = question_id
        if "问题" in group.columns:
            question_text = group["问题"].iloc[0]
        
        eval_data[question_id] = {
            "question": question_text,
            "passages": passages,
            "gold_scores": gold_scores
        }
    
    # 保存评估数据JSON
    eval_json_path = os.path.join(output_dir, "eval_data.json")
    with open(eval_json_path, 'w', encoding='utf-8-sig') as f:
        json.dump(eval_data, f, ensure_ascii=False, indent=2)
    
    print(f"评估数据已保存到 {eval_json_path}")
    
    # 2. 准备BGE重排模型结果JSON
    bge_results = {}
    
    # 按问题ID分组
    grouped = rerank_df.groupby(question_id_column)
    for question_id, group in grouped:
        # 按知识块索引排序
        group = group.sort_values("知识块索引")
        
        bge_scores = group["BGE重排分数"].tolist()
        
        # 生成排名列表：根据BGE重排分数从高到低排序
        # 创建(索引, 分数)对，然后按分数降序排序
        indexed_scores = [(i, score) for i, score in enumerate(bge_scores)]
        indexed_scores.sort(key=lambda x: x[1], reverse=True)
        
        # 提取排序后的索引列表（按BGE重排分数从高到低）
        bge_rankings = [idx for idx, score in indexed_scores]
        
        bge_results[question_id] = {
            "scores": bge_scores,
            "rankings": bge_rankings
        }
    
    # 保存BGE结果JSON
    bge_json_path = os.path.join(output_dir, "bge_rankings.json")
    with open(bge_json_path, 'w', encoding='utf-8-sig') as f:
        json.dump(bge_results, f, ensure_ascii=False, indent=2)
    
    print(f"BGE重排结果已保存到 {bge_json_path}")
    
    # 3. 准备Qwen重排模型结果JSON
    qwen_results = {}
    
    # 按问题ID分组
    grouped = rerank_df.groupby(question_id_column)
    for question_id, group in grouped:
        # 按知识块索引排序
        group = group.sort_values("知识块索引")
        
        qwen_scores = group["Qwen重排分数"].tolist()
        
        # 生成排名列表：根据Qwen重排分数从高到低排序
        # 创建(索引, 分数)对，然后按分数降序排序
        indexed_scores = [(i, score) for i, score in enumerate(qwen_scores)]
        indexed_scores.sort(key=lambda x: x[1], reverse=True)
        
        # 提取排序后的索引列表（按Qwen重排分数从高到低）
        qwen_rankings = [idx for idx, score in indexed_scores]
        
        qwen_results[question_id] = {
            "scores": qwen_scores,
            "rankings": qwen_rankings
        }
    
    # 保存Qwen结果JSON
    qwen_json_path = os.path.join(output_dir, "qwen_rankings.json")
    with open(qwen_json_path, 'w', encoding='utf-8-sig') as f:
        json.dump(qwen_results, f, ensure_ascii=False, indent=2)
    
    print(f"Qwen重排结果已保存到 {qwen_json_path}")
    
    # 4. 准备Bad Case列表JSON (可选)
    # 这里简单地选择前3个问题作为Bad Case示例
    bad_cases = list(eval_data.keys())[:min(3, len(eval_data))]
    
    bad_cases_json_path = os.path.join(output_dir, "bad_cases.json")
    with open(bad_cases_json_path, 'w', encoding='utf-8-sig') as f:
        json.dump(bad_cases, f, ensure_ascii=False, indent=2)
    
    print(f"Bad Case列表已保存到 {bad_cases_json_path}")
    
    return {
        "eval_data": eval_json_path,
        "bge_rankings": bge_json_path,
        "qwen_rankings": qwen_json_path,
        "bad_cases": bad_cases_json_path
    }

def run_evaluation_tool(eval_data_path: str, 
                       bge_rankings_path: str, 
                       qwen_rankings_path: str, 
                       output_dir: str,
                       bad_cases_path: Optional[str] = None):
    """
    调用评估工具进行评估
    
    参数:
        eval_data_path: 评估数据JSON文件路径
        bge_rankings_path: BGE重排结果JSON文件路径
        qwen_rankings_path: Qwen重排结果JSON文件路径
        output_dir: 输出目录
        bad_cases_path: Bad Case列表JSON文件路径（可选）
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 构建命令
    cmd = [
        "python", "-m", "src.main",
        "--data", eval_data_path,
        "--bge-scores", bge_rankings_path,
        "--qwen-scores", qwen_rankings_path,
        "--output-dir", output_dir
    ]
    
    # 如果提供了Bad Case列表，添加到命令中
    if bad_cases_path:
        cmd.extend(["--bad-cases", bad_cases_path])
    
    # 执行命令
    print(f"执行评估命令: {' '.join(cmd)}")
    subprocess.run(cmd)

def run_integration_pipeline(evaluation_csv: str, 
                           reranking_csv: str, 
                           output_dir: str = "results"):
    """
    运行完整的数据整合和评估流程
    
    参数:
        evaluation_csv: 评估结果CSV文件路径
        reranking_csv: 重排结果CSV文件路径
        output_dir: 输出目录
    """
    # 1. 转换CSV为JSON
    json_files = convert_csv_to_json(
        evaluation_csv=evaluation_csv,
        reranking_csv=reranking_csv,
        output_dir="data/processed"
    )
    
    # 2. 调用评估工具
    run_evaluation_tool(
        eval_data_path=json_files["eval_data"],
        bge_rankings_path=json_files["bge_rankings"],
        qwen_rankings_path=json_files["qwen_rankings"],
        output_dir=output_dir,
        bad_cases_path=json_files["bad_cases"]
    )
    
    print(f"数据整合和评估流程完成，结果已保存到 {output_dir}")

def run_from_python_api(evaluation_csv: str, 
                       reranking_csv: str, 
                       output_dir: str = "results"):
    """
    从Python API调用评估流程
    
    参数:
        evaluation_csv: 评估结果CSV文件路径
        reranking_csv: 重排结果CSV文件路径
        output_dir: 输出目录
    """
    # 导入评估模块
    from src.main import main
    
    # 1. 转换CSV为JSON
    json_files = convert_csv_to_json(
        evaluation_csv=evaluation_csv,
        reranking_csv=reranking_csv,
        output_dir="data/processed"
    )
    
    # 2. 构建参数
    import sys
    sys.argv = [
        "src.main",
        "--data", json_files["eval_data"],
        "--bge-scores", json_files["bge_rankings"],
        "--qwen-scores", json_files["qwen_rankings"],
        "--output-dir", output_dir,
        "--bad-cases", json_files["bad_cases"]
    ]
    
    # 3. 调用评估函数
    main()
    
    print(f"数据整合和评估流程完成，结果已保存到 {output_dir}")

if __name__ == "__main__":
    # 示例用法
    evaluation_csv = "data/processed/evaluation_results.csv"
    reranking_csv = "data/processed/reranking_results.csv"
    output_dir = "results"
    
    # 使用命令行方式调用
    run_integration_pipeline(evaluation_csv, reranking_csv, output_dir)
    
    # 或者使用Python API方式调用
    # run_from_python_api(evaluation_csv, reranking_csv, output_dir) 